#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import pandas as pd
from PyQt5.QtWidgets import (QTableView, QFileDialog, QMessageBox,
                             QHeaderView, QVBoxLayout)
from PyQt5.QtCore import Qt, QSortFilterProxyModel
from PyQt5.QtGui import (QStandardItemModel, QStandardItem)
from utils.huitu3d import ThreeDPlotWidget


class DataImportTableView:
    """数据导入与表格展示模块"""

    def __init__(self, ui):
        self.ui = ui
        # 数据模型
        self.table_model = None
        self.search_proxy_model = None
        # 存储原始数据的DataFrame
        self.original_data = None
        # 3D绘图相关
        self.three_d_plot_widget = None
        self.use_dogleg_severity = True  # 默认使用狗腿度
        # 初始化UI
        self.init_ui()

    def init_ui(self):
        """初始化用户界面"""
        # 导入按钮
        self.ui.import_btn.clicked.connect(self.import_data)
        self.ui.import_btn.setToolTip("从Excel或TXT文件导入数据")
        # 连接绘图交互按钮
        if hasattr(self.ui, 'pushButton_6'):
            # 放大
            self.ui.pushButton_6.clicked.connect(self.zoom_in_chart)
        if hasattr(self.ui, 'pushButton_7'):
            # 缩小
            self.ui.pushButton_7.clicked.connect(self.zoom_out_chart)
        if hasattr(self.ui, 'pushButton_5'):
            # 重置
            self.ui.pushButton_5.clicked.connect(self.reset_chart_view)
        # 创建表格视图
        self.ui.tableView.setSelectionBehavior(QTableView.SelectRows)
        self.ui.tableView.setSelectionMode(QTableView.ExtendedSelection)
        self.ui.tableView.setSortingEnabled(False)  # 禁用排序功能
        self.ui.tableView.horizontalHeader().setSectionResizeMode(QHeaderView.Interactive)
        self.ui.tableView.horizontalHeader().setStretchLastSection(True)
        self.ui.tableView.setAlternatingRowColors(True)  # 启用交替行颜色
        self.ui.tableView.setShowGrid(True)  # 显示网格
        self.ui.tableView.setGridStyle(Qt.SolidLine)
        self.ui.tableView.setEditTriggers(QTableView.NoEditTriggers)  # 禁止直接编辑
        header = self.ui.tableView.horizontalHeader()
        header.setSectionResizeMode(QHeaderView.Stretch)
        # 确保plot_widget有布局
        if not self.ui.Draw_widget.layout():
            self.ui.Draw_widget.setLayout(QVBoxLayout())
        # 初始化当前图表视图引用
        self.current_chart_view = None
        # 初始化表格模型
        self.init_models()
        self.ui.plot_btn.clicked.connect(self.show_plot_dialog)
        self.ui.plot_btn.setToolTip("根据导入的数据绘制图表")
        self.ui.plot_btn.setEnabled(False)  # 初始禁用
        # 自动加载demo.txt文件
        self.auto_load_demo_file()

    def init_models(self):
        """初始化表格模型"""
        # 创建原始数据模型
        self.table_model = QStandardItemModel()

        # 创建搜索代理模型
        self.search_proxy_model = QSortFilterProxyModel()
        self.search_proxy_model.setFilterCaseSensitivity(Qt.CaseInsensitive)
        self.search_proxy_model.setSourceModel(self.table_model)

    def import_data(self):
        """从文件导入数据"""
        # 选择导入文件
        file_path, _ = QFileDialog.getOpenFileName(
            None,  # 使用None作为父窗口，避免类型错误
            "选择导入文件",
            "",
            "Excel文件 (*.xlsx *.xls);;CSV文件 (*.csv);;文本文件 (*.txt);;所有文件 (*)"
        )
        if not file_path:
            return
        try:
            # 直接调用load_file方法加载文件
            self.load_file(file_path)
            # 导入成功提示
            QMessageBox.information(None, "导入成功", f"成功导入数据\n请点击'绘图'按钮选择图表类型")
        except Exception as e:
            QMessageBox.warning(None, "导入失败", f"导入数据时发生错误: {str(e)}")

    def validate_data(self, dataframe):
        """验证数据是否符合要求"""
        # 必需的列名（井斜深、井垂深、井斜角、方位角）
        required_columns = ['井斜深', '井垂深', '井斜角', '方位角']
        # 检查必需的列是否存在
        missing_columns = []
        for col in required_columns:
            # 尝试使用精确匹配
            if col not in dataframe.columns:
                # 尝试使用模糊匹配（不区分大小写）
                matched = False
                for df_col in dataframe.columns:
                    if isinstance(df_col, str) and df_col.lower() == col.lower():
                        matched = True
                        break
                if not matched:
                    missing_columns.append(col)
        if missing_columns:
            QMessageBox.warning(None, "数据缺失", f"文件中缺少以下必需列：{', '.join(missing_columns)}")
            return False
        # 检查必需列中是否有空白数据
        empty_cells = []
        for col in required_columns:
            # 查找实际的列名（不区分大小写）
            actual_col = None
            for df_col in dataframe.columns:
                if isinstance(df_col, str) and df_col.lower() == col.lower():
                    actual_col = df_col
                    break
            if actual_col:
                # 检查是否有空值
                if dataframe[actual_col].isna().any():
                    # 找到所有空白行
                    empty_rows = dataframe[dataframe[actual_col].isna()].index.tolist()
                    # 转换为1-based行号
                    empty_rows = [str(row + 2) for row in empty_rows]  # +2 因为索引从0开始，表头占一行
                    empty_cells.append(f"'{actual_col}'列（行：{', '.join(empty_rows)}）")

        if empty_cells:
            QMessageBox.warning(None, "数据不完整", f"以下必需列中存在空白数据：\n{', '.join(empty_cells)}")
            return False

        return True

    def load_data_to_model(self, dataframe):
        """将DataFrame数据加载到表格模型"""
        # 清空现有数据
        self.table_model.clear()
        # 设置表头
        self.table_model.setHorizontalHeaderLabels(dataframe.columns.tolist())
        # 添加数据行
        for _, row in dataframe.iterrows():
            items = [QStandardItem(str(value) if pd.notna(value) else "") for value in row]
            self.table_model.appendRow(items)
        # 将搜索代理模型设置到表格视图
        self.ui.tableView.setModel(self.search_proxy_model)

    def show_plot_dialog(self):
        self.three_d_plot_widget = ThreeDPlotWidget(self.original_data, self.ui.Draw_widget)
        self.three_d_plot_widget.show_plot_dialog(self.ui.Trajectory_Center_Widget)
        # self.three_d_plot_widget.set_stuck_point(100)

    def zoom_in_chart(self):
        # 放大图表
        if self.three_d_plot_widget and self.three_d_plot_widget.isVisible():
            self.three_d_plot_widget.zoom_in()

    def zoom_out_chart(self):
        # 缩小图表
        if self.three_d_plot_widget and self.three_d_plot_widget.isVisible():
            self.three_d_plot_widget.zoom_out()

    def reset_chart_view(self):
        # 重置图表视图到初始状态
        if self.three_d_plot_widget and self.three_d_plot_widget.isVisible():
            self.three_d_plot_widget.reset_view()

    def auto_load_demo_file(self):
        """自动加载demo.txt文件"""
        # 构建demo.txt文件的绝对路径
        import os
        current_dir = os.path.dirname(os.path.abspath(__file__))
        # 向上两级到项目根目录，然后进入data目录
        demo_file_path = os.path.join(current_dir, "..", "data", "demo.txt")
        demo_file_path = os.path.normpath(demo_file_path)

        # 检查文件是否存在
        if os.path.exists(demo_file_path):
            try:
                # 直接调用load_file方法加载文件
                self.load_file(demo_file_path)
            except Exception as e:
                print(f"自动加载demo.txt失败: {str(e)}")

    def load_file(self, file_path):
        """加载指定路径的文件"""
        try:
            # 清除当前绘制的图片
            if self.three_d_plot_widget:
                self.three_d_plot_widget.clear_plot()
            # 根据文件扩展名选择读取方法
            file_ext = os.path.splitext(file_path)[1].lower()
            if file_ext in ['.xlsx', '.xls']:
                # 读取Excel文件
                df = pd.read_excel(file_path)
            elif file_ext == '.csv':
                # 读取CSV文件
                df = pd.read_csv(file_path, encoding='utf-8')
            elif file_ext == '.txt':
                # 读取文本文件
                df = pd.read_csv(file_path, sep='\t', encoding='utf-8')
            else:
                raise ValueError(f"不支持的文件格式: {file_ext}")
            # 数据验证
            if not self.validate_data(df):
                return
            # 加载数据到表格模型
            self.load_data_to_model(df)
            # 保存原始数据
            self.original_data = df.copy()
            # 启用绘图按钮
            self.ui.plot_btn.setEnabled(True)
        except Exception as e:
            raise e
